package BRU

import BRU.Bru
import chisel3._
import chisel3.util._
import BRU.BruInstr._



class Pstate extends Bundle {
  val n = Input(UInt(1.W))
  val z = Input(UInt(1.W))
  val c = Input(UInt(1.W))
  val v = Input(UInt(1.W))
}

class Bru extends Module {

  val io = IO ( new Bundle {
    val pstate   = Input(new Pstate)
    val bpuIn    = Input(new ExStageDataIn)
    val inst     = Input(UInt(32.W))
    val PC       = Input(UInt(64.W))
    val regRn    = Input(UInt(64.W))
    val regRt    = Input(UInt(64.W))
    //    val RnIndex  = Output(UInt(5.W))                //BR、BLR rd_reg
    //    val RtIndex  = Output(UInt(5.W))                //CBNZ、CBN、TBNZ、TBZ rd_reg
    //    val RdIndex  = Output(UInt(5.W))                //ADR、ADRP
    val wbFlag = Output(Bool())
    val wbData = Output(UInt(64.W))
    val brout    = Output(new BrDataIO)

  })

  val BHit     = Wire(Bool())
  val BLHit    = Wire(Bool())
  val BRHit    = Wire(Bool())
  val BLRHit   = Wire(Bool())
  val RETHit   = Wire(Bool())
  val BCONHit  = Wire(Bool())
  val CBNZHit  = Wire(Bool())
  val CBZHit   = Wire(Bool())
  val TBNZHit  = Wire(Bool())
  val TBZHit   = Wire(Bool())
  val ADRHit   = Wire(Bool())
  val ADRPHit  = Wire(Bool())

  BHit    := Mux(io.inst === B, 1.U, 0.U)
  BLHit   := Mux(io.inst === BL, 1.U, 0.U)
  BRHit   := Mux(io.inst === BR, 1.U, 0.U)
  BLRHit  := Mux(io.inst === BLR, 1.U, 0.U)
  RETHit  := Mux(io.inst === RET, 1.U, 0.U)
  BCONHit := Mux(io.inst === BCON, 1.U, 0.U)
  CBNZHit := Mux(io.inst === CBNZ, 1.U, 0.U)
  CBZHit  := Mux(io.inst === CBZ, 1.U, 0.U)
  TBNZHit := Mux(io.inst === TBNZ, 1.U, 0.U)
  TBZHit  := Mux(io.inst === TBZ, 1.U, 0.U)
  ADRHit  := Mux(io.inst === ADR, 1.U, 0.U)
  ADRPHit := Mux(io.inst === ADRP, 1.U, 0.U)

  //generate to rd_reg
  val RnIndex_r  = RegInit(0.U(5.W))
  val RtIndex_r  = RegInit(0.U(5.W))
  RnIndex_r := io.inst(9,5)
  RtIndex_r := io.inst(4,0)
  //  io.RnIndex := RnIndex_r
  //  io.RtIndex := RtIndex_r

  // ADR、ADRP -> 64bit
  val immADR = Wire(UInt(64.W))
  immADR := Cat(Fill(43,io.inst(23)),Cat(io.inst(23,5),io.inst(30,29)))
  val immADRP = Wire(UInt(64.W))
  immADRP := Cat(Fill(31,io.inst(23)),Cat(io.inst(23,5),io.inst(30,29)),Fill(12,0.U))

  // B、BL imm[25:0] -> 64bit
  val imm28_26 = Wire(UInt(28.W))
  imm28_26 := Cat(io.inst(25,0),Fill(2,0.U))
  val imm64_28 = Wire(UInt(64.W))
  imm64_28 := Cat(Fill(36,imm28_26(27)),imm28_26)

  // B.con、CBZ、CBNZ imm[18:0] -> 64bit
  val imm21_19 = Wire(UInt(21.W))
  imm21_19 := Cat(io.inst(23,5),Fill(2,0.U))
  val imm64_21 = Wire(UInt(64.W))
  imm64_21 := Cat(Fill(43,imm21_19(20)),imm21_19)

  // TBNZ_TYPE、TBZ_TYPE imm[13:0] -> 64bit
  val imm16_14 = Wire(UInt(16.W))
  imm16_14 := Cat(io.inst(18,5),Fill(2,0.U))
  val imm64_16 = Wire(UInt(64.W))
  imm64_16 := Cat(Fill(49,imm16_14(15)),imm16_14)

  // BcondFlag generate
  val BcondFlag = Wire(Bool())
  BcondFlag := MuxCase(0.U(1.W),
    Seq(
      (io.inst(3,0) === "b0000".U  &&  io.pstate.z === 1.U) -> 1.U,  //BEQ
      (io.inst(3,0) === "b0001".U  &&  io.pstate.z === 0.U) -> 1.U,  //BNE
      (io.inst(3,0) === "b0010".U  &&  io.pstate.c === 1.U) -> 1.U,  //BCS
      (io.inst(3,0) === "b0011".U  &&  io.pstate.c === 0.U) -> 1.U,  //BCC
      (io.inst(3,0) === "b0100".U  &&  io.pstate.n === 1.U) -> 1.U,  //BMI
      (io.inst(3,0) === "b0101".U  &&  io.pstate.n === 0.U) -> 1.U,  //BPL
      (io.inst(3,0) === "b0110".U  &&  io.pstate.v === 1.U) -> 1.U,  //BVS
      (io.inst(3,0) === "b0111".U  &&  io.pstate.v === 0.U) -> 1.U,  //BVC
      (io.inst(3,0) === "b1000".U  &&  (io.pstate.c === 1.U && io.pstate.z === 0.U)) -> 1.U,          //BHI
      (io.inst(3,0) === "b1001".U  &&  !(io.pstate.c === 1.U && io.pstate.z === 0.U)) -> 1.U,         //BLS
      (io.inst(3,0) === "b1010".U  &&  (io.pstate.n === io.pstate.v)) -> 1.U,                         //BGE
      (io.inst(3,0) === "b1011".U  &&  !(io.pstate.n === io.pstate.v)) -> 1.U,                        //BLT
      (io.inst(3,0) === "b1100".U  &&  (io.pstate.n === io.pstate.v && io.pstate.z === 0.U)) -> 1.U,  //BGT
      (io.inst(3,0) === "b1101".U  &&  !(io.pstate.n === io.pstate.v && io.pstate.z === 0.U)) -> 1.U, //BLE
      (io.inst(3,0) === "b1110".U) -> 1.U,                                                            //AL
      (io.inst(3,0) === "b1111".U) -> 1.U                                                             //AL
    )
  )

  // CBZ/CBNZ Flag generate
  val C32Flag      = Wire(Bool())
  val CBZJumpFlag  = Wire(Bool())
  val CBNZJumpFlag = Wire(Bool())
  C32Flag      := Mux(io.inst(31)  ===  0.U, 1.U, 0.U)
  when(C32Flag){
    when(io.regRt(31,0) === 0.U){
      CBZJumpFlag  := 1.U
      CBNZJumpFlag := 0.U
    }.otherwise{
      CBZJumpFlag  := 0.U
      CBNZJumpFlag := 1.U
    }
  }.otherwise{
    when(io.regRt === 0.U){
      CBZJumpFlag  := 1.U
      CBNZJumpFlag := 0.U
    }.otherwise{
      CBZJumpFlag  := 0.U
      CBNZJumpFlag := 1.U
    }
  }

  // TBZ/TBNZ Flag generate
  val TBZJumpFlag  = Wire(Bool())
  val TBNZJumpFlag = Wire(Bool())
  TBZJumpFlag   := Mux(io.regRt(Cat(io.inst(31),io.inst(23,19))) ===  0.U, 1.U, 0.U)
  TBNZJumpFlag  := Mux(io.regRt(Cat(io.inst(31),io.inst(23,19))) ===  1.U, 1.U, 0.U)

  //Pipe Condition Generate
  val BJump_r     = RegInit(false.B)
  val BLJump_r    = RegInit(false.B)
  val BRJump_r    = RegInit(false.B)
  val BLRJump_r   = RegInit(false.B)
  val RETJump_r   = RegInit(false.B)
  val BCONJump_r  = RegInit(false.B)
  val CBNZJump_r  = RegInit(false.B)
  val CBZJump_r   = RegInit(false.B)
  val TBNZJump_r  = RegInit(false.B)
  val TBZJump_r   = RegInit(false.B)

  BJump_r    := BHit
  BLJump_r   := BLHit
  BRJump_r   := BRHit
  BLRJump_r  := BLRHit
  RETJump_r  := RETHit
  BCONJump_r := BCONHit && BcondFlag
  CBNZJump_r := CBNZHit && CBNZJumpFlag
  CBZJump_r  := CBZHit  && CBZJumpFlag
  TBNZJump_r := TBNZHit && TBNZJumpFlag
  TBZJump_r  := TBZHit  && TBZJumpFlag

  //pip inst type

  val BHit_r     = RegInit(false.B)
  val BLHit_r    = RegInit(false.B)
  val BRHit_r    = RegInit(false.B)
  val BLRHit_r   = RegInit(false.B)
  val RETHit_r   = RegInit(false.B)
  val BCONHit_r  = RegInit(false.B)
  val CBNZHit_r  = RegInit(false.B)
  val CBZHit_r   = RegInit(false.B)
  val TBNZHit_r  = RegInit(false.B)
  val TBZHit_r   = RegInit(false.B)

  BHit_r    := BHit
  BLHit_r   := BLHit
  BRHit_r   := BRHit
  BLRHit_r  := BLRHit
  RETHit_r  := RETHit
  BCONHit_r := BCONHit
  CBNZHit_r := CBNZHit
  CBZHit_r  := CBZHit
  TBNZHit_r := TBNZHit
  TBZHit_r  := TBZHit

//  Link operation、ADR、ADRP
  val WbFlag_r  = RegInit(0.U(1.W))
  val WbData_r  = RegInit(0.U(64.W))
  val RdIndex   = RegInit(0.U(5.W))
  when(BLHit === 1.U || BLRHit === 1.U){
    WbFlag_r   := 1.U
    RdIndex    := 30.U
    WbData_r   := io.PC + 4.U
  }.elsewhen(ADRHit === 1.U){
    WbFlag_r  := 1.U
    RdIndex   := io.inst(4,0)
    WbData_r  := io.PC + immADR
  }.elsewhen(ADRPHit === 1.U){
    WbFlag_r  := 1.U
    RdIndex   := io.inst(4,0)
    WbData_r  := io.PC + immADRP
  }.otherwise{
    WbFlag_r  := 0.U
    RdIndex   := 0.U
    WbData_r  := 0.U
  }

  io.wbData := WbData_r
  io.wbFlag := WbFlag_r


  //PC_r IMM_r regRn_r generate
  val PC_r       = RegInit(0.U(64.W))
  val imm64_28_r = RegInit(0.U(64.W))
  val imm64_21_r = RegInit(0.U(64.W))
  val imm64_16_r = RegInit(0.U(64.W))
  val regRn_r    = RegInit(0.U(64.W))
  PC_r          := io.PC
  imm64_28_r    := imm64_28
  imm64_21_r    := imm64_21
  imm64_16_r    := imm64_16
  regRn_r       := io.regRn

  //jump_addr generate
  val jumpAddr = Wire(UInt(64.W))
  jumpAddr := MuxCase(PC_r + 4.U,
    Seq(
      (BJump_r === 1.U    || BLJump_r === 1.U)                            -> (PC_r + imm64_28_r),
      (BCONJump_r === 1.U || CBZJump_r === 1.U ||  CBNZJump_r === 1.U)    -> (PC_r + imm64_21_r),
      (TBNZJump_r === 1.U ||  TBZJump_r === 1.U)                          -> (PC_r + imm64_16_r),
      (BRJump_r === 1.U   ||  BLRJump_r === 1.U  ||   RETJump_r === 1.U)  -> regRn_r
    )
  )

  //taken generate
  val taken = Wire(Bool())
  val cond_taken = Wire(Bool())
  val nocond_taken = Wire(Bool())
  cond_taken := BCONJump_r || CBNZJump_r || CBZJump_r || TBNZJump_r || TBZJump_r
  nocond_taken := BJump_r || BLJump_r || BRJump_r || BLRJump_r || RETJump_r
  taken := cond_taken || nocond_taken

  //brType generate
  val brType = Wire(UInt(4.W))
  brType := MuxCase(0.U(4.W),
    Seq(
      BHit_r    -> 1.U,
      BLHit_r   -> 2.U,
      BRHit_r   -> 3.U,
      BLRHit_r  -> 4.U,
      RETHit_r  -> 5.U,
      BCONHit_r -> 6.U,
      CBNZHit_r -> 7.U,
      CBZHit_r  -> 8.U,
      TBNZHit_r -> 9.U,
      TBZHit_r  -> 10.U
    )
  )

  val condInstFlag = Wire(Bool())
  condInstFlag := BCONHit_r || CBNZHit_r || CBZHit_r || TBNZHit_r || TBZHit_r
  val nocondInstFlag = Wire(Bool())
  nocondInstFlag := BHit_r || BLHit_r || BRHit_r || BLRHit_r || RETHit_r

  //misPredict generate
  val predictTaken = RegInit(false.B)
  val predictAddr  = RegInit(0.U(64.W))
  predictTaken := io.bpuIn.vld
  predictAddr  := io.bpuIn.idm

  val misPredict = Wire(Bool())
  when(condInstFlag){
    when(cond_taken === predictTaken){
      when(cond_taken === 1.U){
        when(jumpAddr === predictAddr){
          misPredict := 0.U
        }.otherwise{
          misPredict := 1.U
        }
      }.otherwise{
        misPredict := 0.U
      }
    }.otherwise{
      misPredict := 1.U
    }
  }.elsewhen(nocondInstFlag){
    when(jumpAddr === predictAddr){
      misPredict := 0.U
    }.otherwise{
      misPredict := 1.U
    }
  }.otherwise{
    misPredict := 0.U
  }

  //brout generate

  io.brout.brTaken    := taken
  io.brout.brTarget   := jumpAddr
  io.brout.brType     := brType
  io.brout.misPredict := misPredict
  io.brout.brVal      := 1.U

}

 object Bru_Gen extends App {
  println("Generating the adder hardware")
  emitVerilog(new Bru(), Array("--target-dir", "generated/Bru/"))
}


